from imblearn.over_sampling import RandomOverSampler
from imblearn.under_sampling import RandomUnderSampler
from sklearn.model_selection import train_test_split

from src.util.common_util import printx, while_input, is_float, is_int_between


def predict_sample(df):
    printx("[训练]开始读取、合并数据")
    x_data = df.drop(columns=['is_dangerous'])
    y = df["is_dangerous"]
    printx("[训练]划分训练集和测试集")
    test_size = float(while_input("请输入测试集占比(建议0.2):", is_float))
    X_train, x_test, y_train, y_test = train_test_split(x_data, y, test_size=test_size, random_state=42,
                                                        stratify=y)
    sample_type = while_input("请输入采样类型[1:过采样 其他:欠采样]:", is_int_between, (1, 2))
    if f"{sample_type}" == '1':
        printx("[训练]开始过采样")
        x_resampled, y_resampled = RandomOverSampler(random_state=42).fit_resample(X_train, y_train)
    else:
        printx("[训练]开始欠采样")
        x_resampled, y_resampled = RandomUnderSampler(random_state=42).fit_resample(X_train, y_train)
    return x_resampled, y_resampled, x_test, y_test, test_size, sample_type
